{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code Hierarchy Node Parser\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/llama-index-packs/llama-index-packs-code-hierarchy/examples/CodeHierarchyNodeParserUsage.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "The `CodeHierarchyNodeParser` is useful to split long code files into more reasonable chunks. What this will do is create a \"Hierarchy\" of sorts, where sections of the code are made more reasonable by replacing the scope body with short comments telling the LLM to search for a referenced node if it wants to read that context body. This is called skeletonization, and is toggled by setting `skeleton` to `True` which it is by default. Nodes in this hierarchy will be split based on scope, like function, class, or method scope, and will have links to their children and parents so the LLM can traverse the tree.\n",
    "\n",
    "This notebook gives an initial demo of the pack, and then dives into a deeper technical exploration of how it works.\n",
    "\n",
    "**NOTE:** Currently, this pack is configured to only work with `OpenAI` LLMs. But feel free to copy/download the source code and edit as needed!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation and Import\n",
    "\n",
    "First be sure to install the necessary [tree-sitter](https://tree-sitter.github.io/tree-sitter/) libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: tree-sitter in /home/ryanpeach/.pyenv/versions/3.11.6/envs/llama_index/lib/python3.11/site-packages (0.20.4)\n",
      "Requirement already satisfied: tree-sitter-languages in /home/ryanpeach/.pyenv/versions/3.11.6/envs/llama_index/lib/python3.11/site-packages (1.10.2)\n",
      "Requirement already satisfied: python-dotenv in /home/ryanpeach/.pyenv/versions/3.11.6/envs/llama_index/lib/python3.11/site-packages (1.0.1)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.2.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install llama-index-packs-code-hierarchy llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "from llama_index.core.text_splitter import CodeSplitter\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.packs.code_hierarchy import CodeHierarchyNodeParser\n",
    "from llama_index.packs.code_hierarchy import CodeHierarchyAgentPack\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown, display\n",
    "\n",
    "\n",
    "def print_python(python_text):\n",
    "    \"\"\"This function prints python text in ipynb nicely formatted.\"\"\"\n",
    "    display(Markdown(\"```python\\n\" + python_text + \"```\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial Demo\n",
    "\n",
    "First, lets run the pack by using nodes from the included `CodeHierarchyNodeParser`, and from there, explore further how it actually works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model=\"gpt-4\", temperature=0.2)\n",
    "\n",
    "documents = SimpleDirectoryReader(\n",
    "    input_files=[Path(\"../llama_index/packs/code_hierarchy/code_hierarchy.py\")],\n",
    "    file_metadata=lambda x: {\"filepath\": x},\n",
    ").load_data()\n",
    "\n",
    "split_nodes = CodeHierarchyNodeParser(\n",
    "    language=\"python\",\n",
    "    # You can further parameterize the CodeSplitter to split the code\n",
    "    # into \"chunks\" that match your context window size using\n",
    "    # chunck_lines and max_chars parameters, here we just use the defaults\n",
    "    code_splitter=CodeSplitter(language=\"python\", max_chars=1000, chunk_lines=10),\n",
    ").get_nodes_from_documents(documents)\n",
    "\n",
    "pack = CodeHierarchyAgentPack(split_nodes=split_nodes, llm=llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Added user message to memory: How does the get_code_hierarchy_from_nodes function from the code hierarchy node parser work? Provide specific implementation details.\n",
      "=== Calling Function ===\n",
      "Calling function: code_search with args: {\n",
      "  \"input\": \"get_code_hierarchy_from_nodes\"\n",
      "}\n",
      "Got output: def get_code_hierarchy_from_nodes(\n",
      "        nodes: Sequence[BaseNode],\n",
      "        max_depth: int = -1,\n",
      "    ) -> Tuple[Dict[str, Any], str]:\n",
      "# Code replaced for brevity. See node_id 1b2cbe9a-5846-4110-aaa5-26327110c9ab\n",
      "========================\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: code_search with args: {\n",
      "  \"input\": \"1b2cbe9a-5846-4110-aaa5-26327110c9ab\"\n",
      "}\n",
      "Got output: # Code replaced for brevity. See node_id ce774d77-8687-4ae5-af74-4a990c085362\n",
      "\"\"\"\n",
      "        Creates a code hierarchy appropriate to put into a tool description or context\n",
      "        to make it easier to search for code.\n",
      "\n",
      "        Call after `get_nodes_from_documents` and pass that output to this function.\n",
      "        \"\"\"\n",
      "        out: Dict[str, Any] = defaultdict(dict)\n",
      "\n",
      "        def get_subdict(keys: list[str]) -> Dict[str, Any]:\n",
      "            # Get the dictionary we are operating on\n",
      "                # Code replaced for brevity. See node_id 409c28d7-6fc2-434e-99ad-e870bf22e963\n",
      "\n",
      "        def recur_inclusive_scope(node: BaseNode, i: int, keys: list[str]) -> None:\n",
      "                # Code replaced for brevity. See node_id 0fb89089-8cb3-40ab-af65-d655b7511430\n",
      "\n",
      "        def dict_to_markdown(d: dict[str, Any], depth: int = 0) -> str:\n",
      "                # Code replaced for brevity. See node_id 27ef7ca1-0709-4dfd-9cd0-0e2c6b930fa0\n",
      "# Code replaced for brevity. See node_id e8fe25ee-d9bf-4fc0-adf5-6db6ec3020a0\n",
      "========================\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: code_search with args: {\n",
      "  \"input\": \"409c28d7-6fc2-434e-99ad-e870bf22e963\"\n",
      "}\n",
      "Got output: def get_subdict(keys: list[str]) -> Dict[str, Any]:\n",
      "            # Get the dictionary we are operating on\n",
      "            this_dict = out\n",
      "            for key in keys:\n",
      "                if key not in this_dict:\n",
      "                    this_dict[key] = defaultdict(dict)\n",
      "                this_dict = this_dict[key]\n",
      "            return this_dict\n",
      "========================\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: code_search with args: {\n",
      "  \"input\": \"0fb89089-8cb3-40ab-af65-d655b7511430\"\n",
      "}\n",
      "Got output: def recur_inclusive_scope(node: BaseNode, i: int, keys: list[str]) -> None:\n",
      "            if \"inclusive_scopes\" not in node.metadata:\n",
      "                raise KeyError(\"inclusive_scopes not in node.metadata\")\n",
      "            if i >= len(node.metadata[\"inclusive_scopes\"]):\n",
      "                return\n",
      "            scope = node.metadata[\"inclusive_scopes\"][i]\n",
      "\n",
      "            this_dict = get_subdict(keys)\n",
      "\n",
      "            if scope[\"name\"] not in this_dict:\n",
      "                this_dict[scope[\"name\"]] = defaultdict(dict)\n",
      "\n",
      "            if i < max_depth or max_depth == -1:\n",
      "                recur_inclusive_scope(node, i + 1, [*keys, scope[\"name\"]])\n",
      "========================\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: code_search with args: {\n",
      "  \"input\": \"27ef7ca1-0709-4dfd-9cd0-0e2c6b930fa0\"\n",
      "}\n",
      "Got output: def dict_to_markdown(d: dict[str, Any], depth: int = 0) -> str:\n",
      "            markdown = \"\"\n",
      "            indent = \"  \" * depth  # Two spaces per depth level\n",
      "\n",
      "            for key, value in d.items():\n",
      "                if isinstance(value, dict):  # Check if value is a dict\n",
      "                    # Add the key with a bullet point and increase depth for nested dicts\n",
      "                    markdown += f\"{indent}- {key}\\n{dict_to_markdown(value, depth + 1)}\"\n",
      "                else:\n",
      "                    # Handle non-dict items if necessary\n",
      "                    markdown += f\"{indent}- {key}: {value}\\n\"\n",
      "\n",
      "            return markdown\n",
      "========================\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: code_search with args: {\n",
      "  \"input\": \"e8fe25ee-d9bf-4fc0-adf5-6db6ec3020a0\"\n",
      "}\n",
      "Got output: # Code replaced for brevity. See node_id 1b2cbe9a-5846-4110-aaa5-26327110c9ab\n",
      "for node in nodes:\n",
      "            filepath = node.metadata[\"filepath\"].split(\"/\")\n",
      "            filepath[-1] = filepath[-1].split(\".\")[0]\n",
      "            recur_inclusive_scope(node, 0, filepath)\n",
      "\n",
      "        return out, dict_to_markdown(out)\n",
      "========================\n",
      "\n",
      "The `get_code_hierarchy_from_nodes` function in the `CodeHierarchyNodeParser` class is used to create a code hierarchy that can be put into a tool description or context to make it easier to search for code. This function should be called after `get_nodes_from_documents` and the output of that function should be passed to this function.\n",
      "\n",
      "Here's a breakdown of how this function works:\n",
      "\n",
      "1. It initializes an empty dictionary `out`.\n",
      "\n",
      "2. It defines a helper function `get_subdict(keys)`, which navigates through the `out` dictionary using the provided keys, creating new dictionaries along the way if a key does not exist. It returns the final dictionary it navigates to.\n",
      "\n",
      "3. It defines another helper function `recur_inclusive_scope(node, i, keys)`, which recursively navigates through the inclusive scopes of a node and adds them to the `out` dictionary using `get_subdict(keys)`. If the maximum depth is reached or there are no more inclusive scopes, it stops the recursion.\n",
      "\n",
      "4. It defines a third helper function `dict_to_markdown(d, depth)`, which converts the `out` dictionary to a markdown string. It does this by iterating through the dictionary and adding each key-value pair to the string. If a value is a dictionary, it recursively calls `dict_to_markdown` on that dictionary and increases the depth.\n",
      "\n",
      "5. It then iterates over the provided nodes, splits the filepath of each node into a list of directories, and calls `recur_inclusive_scope` on each node with the filepath as the keys.\n",
      "\n",
      "6. Finally, it returns the `out` dictionary and the markdown string generated by `dict_to_markdown(out)`.\n",
      "\n",
      "This function is a part of the `CodeHierarchyNodeParser` class, which is used to parse code hierarchy nodes.\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    pack.run(\n",
    "        \"How does the get_code_hierarchy_from_nodes function from the code hierarchy node parser work? Provide specific implementation details.\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the agent explored the hierarchy of the code by requesting specific function names and IDs, in order to provide a full explanation of how the function works!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Technical Explanations/Exploration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare your Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Choose a directory you want to scan, and glob for all the code files you want to import.\n",
    "\n",
    "In this case I'm going to glob all \"*.py\" files in the `llama_index/node_parser` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = SimpleDirectoryReader(\n",
    "    input_files=[Path(\"../llama_index/packs/code_hierarchy/code_hierarchy.py\")],\n",
    "    file_metadata=lambda x: {\"filepath\": x},\n",
    ").load_data()\n",
    "\n",
    "split_nodes = CodeHierarchyNodeParser(\n",
    "    language=\"python\",\n",
    "    # You can further parameterize the CodeSplitter to split the code\n",
    "    # into \"chunks\" that match your context window size using\n",
    "    # chunck_lines and max_chars parameters, here we just use the defaults\n",
    "    code_splitter=CodeSplitter(language=\"python\", max_chars=1000, chunk_lines=10),\n",
    ").get_nodes_from_documents(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This should be the code hierarchy node parser itself. Lets have it parse itself!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length of text: 33375\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "from collections import defaultdict\n",
       "from enum import Enum\n",
       "from tree_sitter import Node\n",
       "from typing import Any, Dict, List, Optional, Sequence, Tuple\n",
       "\n",
       "\n",
       "from llama_index.core.bridge.pydantic import BaseModel, Field\n",
       "from llama_index.core.callbacks.base import CallbackManager\n",
       "from llama_index.core.extractors.metadata_extractors import BaseExtractor\n",
       "from llama_index.core.node_parser.interface import NodeParser\n",
       "from llama_index.core.schema import BaseNode, NodeRelationship, TextNode\n",
       "from llama_index.core.text_splitter import CodeSplitter\n",
       "from llama_index.core.utils import get_tqdm_iterable\n",
       "\n",
       "\n",
       "class _SignatureCaptureType(BaseModel):\n",
       "    \"\"\"\n",
       "    Unfortunately some languages need special options for how to make a signature.\n",
       "\n",
       "    For example, html element signatures should include their closing >, there is no\n",
       "    easy way to include this using an always-exclusive system.\n",
       "\n",
       "    However, using an always-inclusive system, python decorators don't work,\n",
       "    as there isn't an easy to define terminator for decorators that is inclusive\n",
       "    to their signature.\n",
       "    \"\"\"\n",
       "\n",
       "    type: str = Field(description=\"The type string to match on.\")\n",
       "    inclusive: bool = Field(\n",
       "        description=(\n",
       "            \"Whether to include the text of the node matched by this type or not.\"\n",
       "        ),\n",
       "    )\n",
       "\n",
       "\n",
       "class _SignatureCaptureOptions(BaseModel):\n",
       "    \"\"\"\n",
       "    Options for capturing the signature of a node.\n",
       "    \"\"\"\n",
       "\n",
       "    start_signature_types: Optional[List[_SignatureCaptureType]] = Field(\n",
       "        None,\n",
       "        descripti\n",
       "\n",
       "# ...```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(f\"Length of text: {len(documents[0].text)}\")\n",
    "print_python(documents[0].text[:1500] + \"\\n\\n# ...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is way too long to fit into the context of our LLM. So what are we to do? Well we will split it. We are going to use the `CodeHierarchyNodeParser` to split the nodes into more reasonable chunks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of nodes after splitting: 90\n"
     ]
    }
   ],
   "source": [
    "split_nodes = CodeHierarchyNodeParser(\n",
    "    language=\"python\",\n",
    "    # You can further parameterize the CodeSplitter to split the code\n",
    "    # into \"chunks\" that match your context window size using\n",
    "    # chunck_lines and max_chars parameters, here we just use the defaults\n",
    "    code_splitter=CodeSplitter(language=\"python\", max_chars=1000, chunk_lines=10),\n",
    ").get_nodes_from_documents(documents)\n",
    "print(\"Number of nodes after splitting:\", len(split_nodes))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! So that split up our data from 1 node into quite a few nodes! Whats the max length of any of these nodes?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Longest text in nodes: 1152\n"
     ]
    }
   ],
   "source": [
    "print(f\"Longest text in nodes: {max(len(n.text) for n in split_nodes)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's much shorter than before! Let's look at a sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "from collections import defaultdict\n",
       "from enum import Enum\n",
       "from tree_sitter import Node\n",
       "from typing import Any, Dict, List, Optional, Sequence, Tuple\n",
       "\n",
       "\n",
       "from llama_index.core.bridge.pydantic import BaseModel, Field\n",
       "from llama_index.core.callbacks.base import CallbackManager\n",
       "from llama_index.core.extractors.metadata_extractors import BaseExtractor\n",
       "from llama_index.core.node_parser.interface import NodeParser\n",
       "from llama_index.core.schema import BaseNode, NodeRelationship, TextNode\n",
       "from llama_index.core.text_splitter import CodeSplitter\n",
       "from llama_index.core.utils import get_tqdm_iterable\n",
       "\n",
       "\n",
       "class _SignatureCaptureType(BaseModel):\n",
       "    # Code replaced for brevity. See node_id b30b6043-4cba-420e-bd6b-e91beea08819\n",
       "\n",
       "\n",
       "class _SignatureCaptureOptions(BaseModel):\n",
       "    # Code replaced for brevity. See node_id e0961aad-bd9f-4295-927d-90ac6e2b06c8\n",
       "# Code replaced for brevity. See node_id 0f6bc262-ef8b-4051-8c8e-486863e4cbe2```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print_python(split_nodes[0].text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Without even needing a long printout we can see everything this module imported in the first document (which is at the module level) and some classes it defines.\n",
    "\n",
    "We also see that it has put comments in place of code that was removed to make the text size more reasonable.\n",
    "These can appear at the beginning or end of a chunk, or at a new scope level, like a class or function declaration.\n",
    "\n",
    "`# Code replaced for brevity. See node_id {node_id}`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Code Hierarchy\n",
    "\n",
    "These scopes can be listed by the `CodeHierarchyNodeParser`, giving a \"repo map\" of sorts.\n",
    "The namesake of this node parser, it creates a tree of scope names to use to search the code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(defaultdict(<class 'dict'>, {'..': defaultdict(<class 'dict'>, {'llama_index': defaultdict(<class 'dict'>, {'packs': defaultdict(<class 'dict'>, {'code_hierarchy': defaultdict(<class 'dict'>, {'code_hierarchy': defaultdict(<class 'dict'>, {'_SignatureCaptureType': defaultdict(<class 'dict'>, {}), '_SignatureCaptureOptions': defaultdict(<class 'dict'>, {}), '_ScopeMethod': defaultdict(<class 'dict'>, {}), '_CommentOptions': defaultdict(<class 'dict'>, {}), '_ScopeItem': defaultdict(<class 'dict'>, {}), '_ChunkNodeOutput': defaultdict(<class 'dict'>, {}), 'CodeHierarchyNodeParser': defaultdict(<class 'dict'>, {'class_name': defaultdict(<class 'dict'>, {}), '__init__': defaultdict(<class 'dict'>, {}), '_get_node_name': defaultdict(<class 'dict'>, {'recur': defaultdict(<class 'dict'>, {})}), '_get_node_signature': defaultdict(<class 'dict'>, {'find_start': defaultdict(<class 'dict'>, {}), 'find_end': defaultdict(<class 'dict'>, {})}), '_chunk_node': defaultdict(<class 'dict'>, {}), 'get_code_hierarchy_from_nodes': defaultdict(<class 'dict'>, {'get_subdict': defaultdict(<class 'dict'>, {}), 'recur_inclusive_scope': defaultdict(<class 'dict'>, {}), 'dict_to_markdown': defaultdict(<class 'dict'>, {})}), '_parse_nodes': defaultdict(<class 'dict'>, {}), '_get_indentation': defaultdict(<class 'dict'>, {}), '_get_comment_text': defaultdict(<class 'dict'>, {}), '_create_comment_line': defaultdict(<class 'dict'>, {}), '_get_replacement_text': defaultdict(<class 'dict'>, {}), '_skeletonize': defaultdict(<class 'dict'>, {}), '_skeletonize_list': defaultdict(<class 'dict'>, {'recur': defaultdict(<class 'dict'>, {})})})})})})})})}), '- ..\\n  - llama_index\\n    - packs\\n      - code_hierarchy\\n        - code_hierarchy\\n          - _SignatureCaptureType\\n          - _SignatureCaptureOptions\\n          - _ScopeMethod\\n          - _CommentOptions\\n          - _ScopeItem\\n          - _ChunkNodeOutput\\n          - CodeHierarchyNodeParser\\n            - class_name\\n            - __init__\\n            - _get_node_name\\n              - recur\\n            - _get_node_signature\\n              - find_start\\n              - find_end\\n            - _chunk_node\\n            - get_code_hierarchy_from_nodes\\n              - get_subdict\\n              - recur_inclusive_scope\\n              - dict_to_markdown\\n            - _parse_nodes\\n            - _get_indentation\\n            - _get_comment_text\\n            - _create_comment_line\\n            - _get_replacement_text\\n            - _skeletonize\\n            - _skeletonize_list\\n              - recur\\n')\n"
     ]
    }
   ],
   "source": [
    "print(CodeHierarchyNodeParser.get_code_hierarchy_from_nodes(split_nodes))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exploration by the Programmer\n",
    "\n",
    "So that we understand what is going on under the hood, what if we go to that node_id we found above?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Going to print the node with UUID: 6d205ded-3ee7-454a-9498-7d5f63963d4c\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "class CodeHierarchyNodeParser(NodeParser):\n",
       "# Code replaced for brevity. See node_id 1b87e4b8-08ef-4b34-ac71-9fbcca8bed76```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "split_nodes_by_id = {n.node_id: n for n in split_nodes}\n",
    "uuid_from_text = split_nodes[9].text.splitlines()[-1].split(\" \")[-1]\n",
    "print(\"Going to print the node with UUID:\", uuid_from_text)\n",
    "print_python(split_nodes_by_id[uuid_from_text].text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the next split in the file. It is prepended with the node before it and appended with the node after it as a comment.\n",
    "\n",
    "We can also see the relationships on this node programmatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{<NodeRelationship.SOURCE: '1'>: RelatedNodeInfo(node_id='6d205ded-3ee7-454a-9498-7d5f63963d4c', node_type=<ObjectType.TEXT: '1'>, metadata={'language': 'python', 'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}], 'start_byte': 6241, 'end_byte': 33374, 'filepath': '../llama_index/packs/code_hierarchy/code_hierarchy.py'}, hash='714b8e8a6c2e99ae5f43521fe600587eda6d2cee8411082c4ba3255701ad443f'),\n",
       " <NodeRelationship.NEXT: '3'>: RelatedNodeInfo(node_id='1b87e4b8-08ef-4b34-ac71-9fbcca8bed76', node_type=<ObjectType.TEXT: '1'>, metadata={}, hash='443391a4ee2bdb50953f94fe46e8d93c8be044ea84c7cc30efdfd5a2234a3c6f'),\n",
       " <NodeRelationship.CHILD: '5'>: [RelatedNodeInfo(node_id='c81c5ec6-02da-43f1-beab-70cdff2ea7e8', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': 'class_name', 'type': 'function_definition', 'signature': 'def class_name(cls) -> str:'}], 'start_byte': 6443, 'end_byte': 6545}, hash='390461d846c3220bc4f5323b1ca81cc5e6d5c2183ac5b61c78919be82574cfb9'),\n",
       "  RelatedNodeInfo(node_id='d027c1a6-bae5-4c7f-be81-ca2e6fb53d67', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': '__init__', 'type': 'function_definition', 'signature': 'def __init__(\\n        self,\\n        language: str,\\n        skeleton: bool = True,\\n        signature_identifiers: Optional[Dict[str, _SignatureCaptureOptions]] = None,\\n        code_splitter: Optional[CodeSplitter] = None,\\n        callback_manager: Optional[CallbackManager] = None,\\n        metadata_extractor: Optional[BaseExtractor] = None,\\n        chunk_min_characters: int = 80,\\n    ):'}], 'start_byte': 7844, 'end_byte': 8992}, hash='89029f5dbeed6fb78c458c8b71a341bb97b9a76aa614b6f267e58ad3ec387556'),\n",
       "  RelatedNodeInfo(node_id='affd8a83-a1bf-4451-9ce6-128080938210', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': '_get_node_name', 'type': 'function_definition', 'signature': 'def _get_node_name(self, node: Node) -> str:'}], 'start_byte': 8994, 'end_byte': 9541}, hash='af1f5405245e154c3bc13a9b796d6f6c4acf21a03431a3b2d160cf88fb9467f8'),\n",
       "  RelatedNodeInfo(node_id='0ee4146b-58b1-4a93-ac31-c36268b4424c', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': '_get_node_signature', 'type': 'function_definition', 'signature': 'def _get_node_signature(self, text: str, node: Node) -> str:'}], 'start_byte': 9543, 'end_byte': 11135}, hash='9c6d0238ac9f394a428be03647887c2bc23b46bd108c0f8946da56e4742a6bda'),\n",
       "  RelatedNodeInfo(node_id='10ed6803-e6d2-4d58-8dd9-52d60b570320', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': '_chunk_node', 'type': 'function_definition', 'signature': 'def _chunk_node(\\n        self,\\n        parent: Node,\\n        text: str,\\n        _context_list: Optional[List[_ScopeItem]] = None,\\n        _root: bool = True,\\n    ) -> _ChunkNodeOutput:'}], 'start_byte': 11137, 'end_byte': 17216}, hash='3c35fe3a0f12b328eeac49c925b39dffffa3d793cfae20a4413f5766b7ce0c1a'),\n",
       "  RelatedNodeInfo(node_id='8399170e-d4d8-4a9e-8d82-8cecd8cf4afd', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': 'get_code_hierarchy_from_nodes', 'type': 'function_definition', 'signature': 'def get_code_hierarchy_from_nodes(\\n        nodes: Sequence[BaseNode],\\n        max_depth: int = -1,\\n    ) -> Tuple[Dict[str, Any], str]:'}], 'start_byte': 17236, 'end_byte': 19499}, hash='59235d7c8606a4b3bab9a5d404c1c4d2d8bb6e4468cd4ee01cf219881fc8fd8f'),\n",
       "  RelatedNodeInfo(node_id='576ef5c2-5ef9-457d-acea-c61d11c6d4e6', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': '_parse_nodes', 'type': 'function_definition', 'signature': 'def _parse_nodes(\\n        self,\\n        nodes: Sequence[BaseNode],\\n        show_progress: bool = False,\\n        **kwargs: Any,\\n    ) -> List[BaseNode]:'}], 'start_byte': 19501, 'end_byte': 25472}, hash='9c392c57f7a24f2757e790522cbca7cab0ec53559567d3d011ce0f15ac78f170'),\n",
       "  RelatedNodeInfo(node_id='c36a088a-5583-4d20-a6f6-c314d987b4c5', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': '_get_indentation', 'type': 'function_definition', 'signature': 'def _get_indentation(text: str) -> Tuple[str, int, int]:'}], 'start_byte': 25492, 'end_byte': 27817}, hash='78da4361d6dc2cd1d475e522f9036343d580e1072cd6f647bf404c4a6859d7d9'),\n",
       "  RelatedNodeInfo(node_id='41c2f0ab-3ab3-4691-acf7-3c6f2162d10a', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': '_get_comment_text', 'type': 'function_definition', 'signature': 'def _get_comment_text(node: TextNode) -> str:'}], 'start_byte': 27837, 'end_byte': 28035}, hash='780851f89f59302c04bbbe13cb46dfa404e04ad7a3e568a4527e1c0a5dbb2812'),\n",
       "  RelatedNodeInfo(node_id='99b4bfca-f48a-4cd7-aa73-fcaeb4fd4461', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': '_create_comment_line', 'type': 'function_definition', 'signature': 'def _create_comment_line(cls, node: TextNode, indention_lvl: int = -1) -> str:'}], 'start_byte': 28054, 'end_byte': 29175}, hash='d5267b1cdc5b60644125f7d196fea602471760a7063e67fb52f5eaadf798a246'),\n",
       "  RelatedNodeInfo(node_id='3f4a4cc7-049d-4bf9-a0f1-9b546508af16', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': '_get_replacement_text', 'type': 'function_definition', 'signature': 'def _get_replacement_text(cls, child_node: TextNode) -> str:'}], 'start_byte': 29194, 'end_byte': 32017}, hash='423740218dc6072242982baf9cc1a3ba5676d9990f80ee41817f7e989aae752c'),\n",
       "  RelatedNodeInfo(node_id='109e1549-e92a-4f76-b2f8-f564ba2fcb46', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': '_skeletonize', 'type': 'function_definition', 'signature': 'def _skeletonize(cls, parent_node: TextNode, child_node: TextNode) -> None:'}], 'start_byte': 32036, 'end_byte': 32702}, hash='cb348537cee003c69c8704dd203aaf0ff93408642ba8c961d52c43168f49c213'),\n",
       "  RelatedNodeInfo(node_id='65505ced-5998-41af-9d22-9b8b6182395f', node_type=<ObjectType.TEXT: '1'>, metadata={'inclusive_scopes': [{'name': 'CodeHierarchyNodeParser', 'type': 'class_definition', 'signature': 'class CodeHierarchyNodeParser(NodeParser):'}, {'name': '_skeletonize_list', 'type': 'function_definition', 'signature': \"def _skeletonize_list(cls, nodes: List[TextNode]) -> None:\\n        # Create a convenient map for mapping node id's to nodes\"}], 'start_byte': 32721, 'end_byte': 33374}, hash='1b60cc9e6e64c851677274b3693e3fdbfd213c4910bf1d2da9a7d502c07d0008')],\n",
       " <NodeRelationship.PARENT: '4'>: RelatedNodeInfo(node_id='1d5bf8bf-639b-486a-adc1-6cd9352eb275', node_type=<ObjectType.TEXT: '1'>, metadata={'language': 'python', 'inclusive_scopes': [], 'start_byte': 0, 'end_byte': 33375, 'filepath': '../llama_index/packs/code_hierarchy/code_hierarchy.py'}, hash='f76ef425d472658be552f10fd2279c0040c2d5ec83a05b7589d97013186e2241')}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "split_nodes_by_id[uuid_from_text].relationships"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `NEXT` `PREV` relationships come from the `CodeSplitter` which is a component of the `CodeHierarchyNodeParser`. It is responsible for cutting up the nodes into chunks that are a certain character length. For more information about the `CodeSplitter` read this:\n",
    "\n",
    "[Code Splitter](https://docs.llamaindex.ai/en/latest/api/llama_index.node_parser.CodeSplitter.html)\n",
    "\n",
    "The `PARENT` and `CHILD` relationships come from the `CodeHierarchyNodeParser` which is responsible for creating the hierarchy of nodes. Things like classes, functions, and methods are nodes in this hierarchy.\n",
    "\n",
    "The `SOURCE` is the original file that this node came from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "# Code replaced for brevity. See node_id 6d205ded-3ee7-454a-9498-7d5f63963d4c\n",
       "\"\"\"Split code using a AST parser.\n",
       "\n",
       "    Add metadata about the scope of the code block and relationships between\n",
       "    code blocks.\n",
       "    \"\"\"\n",
       "\n",
       "    @classmethod\n",
       "    def class_name(cls) -> str:\n",
       "        # Code replaced for brevity. See node_id c81c5ec6-02da-43f1-beab-70cdff2ea7e8\n",
       "\n",
       "    language: str = Field(\n",
       "        description=\"The programming language of the code being split.\"\n",
       "    )\n",
       "    signature_identifiers: Dict[str, _SignatureCaptureOptions] = Field(\n",
       "        description=(\n",
       "            \"A dictionary mapping the type of a split mapped to the first and last type\"\n",
       "            \" of itschildren which identify its signature.\"\n",
       "        )\n",
       "    )\n",
       "    min_characters: int = Field(\n",
       "        default=80,\n",
       "        description=(\n",
       "            \"Minimum number of characters per chunk.Defaults to 80 because that's about\"\n",
       "            \" how long a replacement comment is in skeleton mode.\"\n",
       "        ),\n",
       "    )\n",
       "# Code replaced for brevity. See node_id b5ffc7d6-b795-4304-9dcc-b31568291861```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from llama_index.core.schema import NodeRelationship\n",
    "\n",
    "node_id = uuid_from_text\n",
    "if NodeRelationship.NEXT not in split_nodes_by_id[node_id].relationships:\n",
    "    print(\"No next node found!\")\n",
    "else:\n",
    "    next_node_relationship_info = split_nodes_by_id[node_id].relationships[\n",
    "        NodeRelationship.NEXT\n",
    "    ]\n",
    "    next_node = split_nodes_by_id[next_node_relationship_info.node_id]\n",
    "    print_python(next_node.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keyword Table and Usage by the LLM\n",
    "\n",
    "Lets explore the use of this node parser in an index. We will be able to use any index which allows search by keyword, which should enable us to search for any node by it's uuid, or by any scope name.\n",
    "\n",
    "We have created a `CodeHierarchyKeywordQueryEngine` which will allow us to search for nodes by their uuid, or by their scope name. It's `.query` method can be used as a simple search tool for any LLM. Given the repo map we created earlier, or the text of a split file, the LLM should be able to figure out what to search for very naturally.\n",
    "\n",
    "Lets create the KeywordQueryEngine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.packs.code_hierarchy import CodeHierarchyKeywordQueryEngine\n",
    "\n",
    "query_engine = CodeHierarchyKeywordQueryEngine(\n",
    "    nodes=split_nodes,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can get the same code as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "from collections import defaultdict\n",
       "from enum import Enum\n",
       "from tree_sitter import Node\n",
       "from typing import Any, Dict, List, Optional, Sequence, Tuple\n",
       "\n",
       "\n",
       "from llama_index.core.bridge.pydantic import BaseModel, Field\n",
       "from llama_index.core.callbacks.base import CallbackManager\n",
       "from llama_index.core.extractors.metadata_extractors import BaseExtractor\n",
       "from llama_index.core.node_parser.interface import NodeParser\n",
       "from llama_index.core.schema import BaseNode, NodeRelationship, TextNode\n",
       "from llama_index.core.text_splitter import CodeSplitter\n",
       "from llama_index.core.utils import get_tqdm_iterable\n",
       "\n",
       "\n",
       "class _SignatureCaptureType(BaseModel):\n",
       "    # Code replaced for brevity. See node_id b30b6043-4cba-420e-bd6b-e91beea08819\n",
       "\n",
       "\n",
       "class _SignatureCaptureOptions(BaseModel):\n",
       "    # Code replaced for brevity. See node_id e0961aad-bd9f-4295-927d-90ac6e2b06c8\n",
       "# Code replaced for brevity. See node_id 0f6bc262-ef8b-4051-8c8e-486863e4cbe2```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print_python(query_engine.query(split_nodes[0].node_id).response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But now we can also search for any node by it's common sense name.\n",
    "\n",
    "For example, the class `_SignatureCaptureOptions` is a node in the hierarchy. We can search for it by name.\n",
    "\n",
    "The reason we aren't getting more detail is because our min_characters is too low, try to increase it for more detail for any individual query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "class _SignatureCaptureOptions(BaseModel):\n",
       "# Code replaced for brevity. See node_id f3ccdeee-207a-4d71-9451-7a9aa93bec33```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print_python(query_engine.query(\"_SignatureCaptureOptions\").response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And by module name, in case the LLM sees something in an import statement and wants to know more about it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "from collections import defaultdict\n",
       "from enum import Enum\n",
       "from tree_sitter import Node\n",
       "from typing import Any, Dict, List, Optional, Sequence, Tuple\n",
       "\n",
       "\n",
       "from llama_index.core.bridge.pydantic import BaseModel, Field\n",
       "from llama_index.core.callbacks.base import CallbackManager\n",
       "from llama_index.core.extractors.metadata_extractors import BaseExtractor\n",
       "from llama_index.core.node_parser.interface import NodeParser\n",
       "from llama_index.core.schema import BaseNode, NodeRelationship, TextNode\n",
       "from llama_index.core.text_splitter import CodeSplitter\n",
       "from llama_index.core.utils import get_tqdm_iterable\n",
       "\n",
       "\n",
       "class _SignatureCaptureType(BaseModel):\n",
       "    # Code replaced for brevity. See node_id b30b6043-4cba-420e-bd6b-e91beea08819\n",
       "\n",
       "\n",
       "class _SignatureCaptureOptions(BaseModel):\n",
       "    # Code replaced for brevity. See node_id e0961aad-bd9f-4295-927d-90ac6e2b06c8\n",
       "# Code replaced for brevity. See node_id 0f6bc262-ef8b-4051-8c8e-486863e4cbe2```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print_python(query_engine.query(\"code_hierarchy\").response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### As an Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can convert the query engine to be used as a tool for an agent!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import QueryEngineTool\n",
    "\n",
    "tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=query_engine,\n",
    "    name=\"code_lookup\",\n",
    "    description=\"Useful for looking up information about the code hierarchy codebase.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is also a helpful description of the tool here, which works best as a system prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "Description: Search the tool by any element in this list to get more information about that element.\n",
       "If you see 'Code replaced for brevity' then a uuid, you may also search the tool with that uuid to see the full code.\n",
       "You may need to use the tool multiple times to fully answer the user message.\n",
       "The list is:\n",
       "- ..\n",
       "  - llama_index\n",
       "    - packs\n",
       "      - code_hierarchy\n",
       "        - code_hierarchy\n",
       "          - _SignatureCaptureType\n",
       "          - _SignatureCaptureOptions\n",
       "          - _ScopeMethod\n",
       "          - _CommentOptions\n",
       "          - _ScopeItem\n",
       "          - _ChunkNodeOutput\n",
       "          - CodeHierarchyNodeParser\n",
       "            - class_name\n",
       "            - __init__\n",
       "            - _get_node_name\n",
       "              - recur\n",
       "            - _get_node_signature\n",
       "              - find_start\n",
       "              - find_end\n",
       "            - _chunk_node\n",
       "            - get_code_hierarchy_from_nodes\n",
       "              - get_subdict\n",
       "              - recur_inclusive_scope\n",
       "              - dict_to_markdown\n",
       "            - _parse_nodes\n",
       "            - _get_indentation\n",
       "            - _get_comment_text\n",
       "            - _create_comment_line\n",
       "            - _get_replacement_text\n",
       "            - _skeletonize\n",
       "            - _skeletonize_list\n",
       "              - recur\n",
       "\n"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Markdown(\"Description: \" + query_engine.get_tool_instructions()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets finally actually make an agent!\n",
    "\n",
    "Note that this requires some complex reasoning, and works best with GPT-4-like LLMs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4\", temperature=0.1)\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    tools=[tool],\n",
    "    llm=llm,\n",
    "    system_prompt=query_engine.get_tool_instructions(),\n",
    ")\n",
    "\n",
    "# Create a context for the agent\n",
    "from llama_index.core.workflow import Context\n",
    "\n",
    "ctx = Context(agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await agent.run(\n",
    "    \"How does the get_code_hierarchy_from_nodes function from the code hierarchy node parser work? Provide specific implementation details.\",\n",
    "    ctx=ctx,\n",
    ")\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
